// Deferred Shading Compositor
// (c) jimon game studio

#include "CDeferredShadingCompositor.h"
#include "ISceneManager.h"
#include "IVariables.h"

#define DEFAULT_DS_SHADER \
	"#ifdef VS\n\
	FIN vec4 IPosition;\n\
	FIN vec2 ITexCoord;\n\
	FOUT vec2 TexCoord;\n\
	uniform mat4 MVPMat;\n\
	uniform mat4 MVMat;\n\
	void main()\n\
	{\n\
		TexCoord.x = ITexCoord.x;\n\
		TexCoord.y = 1.0-ITexCoord.y;\n\
		gl_Position = MVPMat * IPosition;\n\
	}\n\
	#endif\n\
	#ifdef FS\n\
	SIN vec2 TexCoord;\n\
	\n\
	uniform float Near;\n\
	uniform float Far;\n\
	uniform float Fov;\n\
	uniform float AspectRatio;\n\
	\n\
	uniform sampler2D Texture;\n\
	uniform sampler2D Texture2;\n\
	uniform sampler2D Texture3;\n\
	\n\
	uniform vec3 PosLights[MAX_DS_BATCH_SIZE];\n\
	uniform int LightsNum;\n\
	\n\
	void main()\n\
	{\n\
		// Get diffuse color\n\
		vec4 DiffuseColor = texture2D(Texture,TexCoord);\n\
	\n\
		// Get normal in view space\n\
		vec3 Normal;\n\
		DS_GET_NORMAL_VALUE(Normal,TexCoord)\n\
	\n\
		// Get position in view space\n\
		float Depth = texture2D(Texture3,TexCoord).z;\n\
		float Z = Far*Near/(Depth*(Near-Far)+Far);\n\
		float TanFovZ = tan(Fov)*Z;\n\
		vec2 UV = (TexCoord*2.0-1.0);\n\
		vec3 Position = vec3(UV.x*TanFovZ*AspectRatio,UV.y*TanFovZ,-Z);\n\
		gl_FragDepth = Depth;\n\
	\n\
		// Result Color\n\
		vec4 ResColor = vec4(0.0);\n\
	\n\
		for(int i = 0; i < LightsNum;i++)\n\
		{\n\
			// Get light pos\n\
			vec3 PosLight = PosLights[i];\n\
		\n\
			// Get light direction\n\
			vec3 LightDir = normalize(PosLight-Position);\n\
		\n\
			// Get diffuse light\n\
			float Diff = max(dot(Normal,LightDir),0.0);\n\
		\n\
			// Get distance light\n\
			float Power = 4;\n\
			float Dist = clamp(1-distance(Position,PosLight)/Power,0.0,1.0);\n\
		\n\
			// Specular\n\
			float Specular = 0.0;\n\
			if(Diff > 0.1)\n\
				Specular = pow(max(-normalize(reflect(LightDir,Normal)).z,0.0),200.0);\n\
		\n\
			// Add to result color\n\
			ResColor += (DiffuseColor*Diff+vec4(Specular))*Dist;\n\
		}\n\
		//ResColor = vec4(DiffuseColor.rgb,0);\n\
	\n\
		// Output color\n\
		//if(DiffuseColor.a > 0.0)\n\
			OUTCOLOR0 = ResColor;\n\
		//else\n\
		//	discard;\n\
	}\n\
	#endif\n"

namespace je
{
	namespace scene
	{
		//! Constructor
		CDeferredShadingCompositor::CDeferredShadingCompositor(ISceneManager * SceneManager)
			:ICompositor(SceneManager,"DeferredShadingCompositor")
		{
			EnableAA = (core::GetVariables()->GetVariablec8("Window_MSAA") > 0) && core::GetVariables()->GetVariablec8("SceneManager_DSCompositor_AA");

			MaxLightsPerPass = core::GetVariables()->GetVariablec8("SceneManager_DSCompositor_Lights_PerPass");

			//Set Base DS Defines
			CurrentSceneManager->GetRender()->SetShaderMaterialDefine("SOLID_VS_VARS","");
			CurrentSceneManager->GetRender()->SetShaderMaterialDefine("SOLID_VS_MAIN","");
			CurrentSceneManager->GetRender()->SetShaderMaterialDefine("SOLID_FS_VARS","USEOUTCOLOR(1)");

			//Set DS Normal Encode/Decode Way
			CurrentSceneManager->GetRender()->SetShaderMaterialDefine("SOLID_FS_MAIN(__Color,__Normal)","OUTCOLOR0 = vec4(__Normal.x,(__Normal.y*0.5+0.5)*sign(__Normal.z),0.0,0.0);\\\nOUTCOLOR1 = __Color;");
			CurrentSceneManager->GetRender()->SetShaderMaterialDefine("DS_GET_NORMAL_VALUE(__Normal,__TexCoord)","vec2 _NormalTempRG = texture2D(Texture2,__TexCoord).xy;\\\nvec2 _NormalTempXY = vec2(_NormalTempRG.x,abs(_NormalTempRG.y)*2.0-1.0);\\\n__Normal = normalize(vec3(_NormalTempXY,sign(_NormalTempRG.y)*sqrt(max(1.0-dot(_NormalTempXY,_NormalTempXY),0.0))));");

			//Set transparent forward render
			CurrentSceneManager->GetRender()->SetShaderMaterialDefine("TRANSPARENT_VS_VARS","FOUT vec3 jeTransparentPosition;\\\nuniform mat4 MVMat;");
			CurrentSceneManager->GetRender()->SetShaderMaterialDefine("TRANSPARENT_VS_MAIN","jeTransparentPosition = (MVMat * IPosition).xyz;");
			CurrentSceneManager->GetRender()->SetShaderMaterialDefine("TRANSPARENT_FS_VARS",jeStringc("SIN vec3 jeTransparentPosition;\\\nuniform vec3 jeForward_TransparentLightPosition[")+jeStringc(MaxLightsPerPass)+jeStringc("];\\\nuniform int jeLightsNumTransparent;"));
			CurrentSceneManager->GetRender()->SetShaderMaterialDefine("TRANSPARENT_FS_MAIN(__Color,__Normal)","\
				vec3 jeTempResult = vec3(0.0);\\\n\
				for(int i = 0; i < jeLightsNumTransparent;i++)\\\n\
				{\\\n\
					vec3 LightDir = normalize(jeForward_TransparentLightPosition[i]-jeTransparentPosition);\\\n\
					float Diff = max(dot(__Normal,LightDir),0.0);\\\n\
					float Power = 4;\\\n\
					float Dist = clamp(1-distance(jeTransparentPosition,jeForward_TransparentLightPosition[i])/Power,0.0,1.0);\\\n\
					float Specular = 0.0;\\\n\
					if(Diff > 0.1)\\\n\
						Specular = pow(max(-normalize(reflect(LightDir,__Normal)).z,0.0),200.0);\\\n\
					jeTempResult += (__Color.rgb*Diff+vec3(Specular))*Dist;\\\n\
				}\\\n\
				OUTCOLOR0 = vec4(jeTempResult,__Color.a);");

			//Create GBuffer FBO
			if(EnableAA)
			{
				FBO_AA = CurrentSceneManager->GetRender()->CreateTextureRenderTarget();
				FBO_Downsample = CurrentSceneManager->GetRender()->CreateTextureRenderTarget();
			}
			FBO_GBuffer = CurrentSceneManager->GetRender()->CreateTextureRenderTarget();

			//Create DS material
			JENEW(DeferredShadingMaterial,video::jeMaterial)
			DeferredShadingMaterial->Type = video::MT_ADDITIVEBLEND;
			DeferredShadingMaterial->EnableZBufferTest = true;
			DeferredShadingMaterial->EnableZBufferMask = true;

			//Create DS Textures
			switch(CreateDSTextures(CurrentSceneManager->GetRender()->GetContext()->GetWindow()->GetSize()))
			{
			case 1:
				{
					//If cant create RG16F, use RGBA8 for normals
					CurrentSceneManager->GetRender()->SetShaderMaterialDefine("SOLID_FS_MAIN(__Color,__Normal)","OUTCOLOR1 = vec2tovec4compress(vec2(__Normal.x*0.5+0.5,(__Normal.y*0.5+0.5)*sign(__Normal.z)*0.5+0.5));");
					CurrentSceneManager->GetRender()->SetShaderMaterialDefine("DS_GET_NORMAL_VALUE(__Normal,__TexCoord)","vec2 _NormalTempRG = vec4tovec2decompress(texture2D(Texture2,__TexCoord));\\\nvec2 _NormalTempXY = vec2(_NormalTempRG.x*2.0-1.0,abs(_NormalTempRG.y*4.0-2.0)-1.0);\\\n__Normal = normalize(vec3(_NormalTempXY,sign(_NormalTempRG.y*2.0-1.0)*sqrt(max(1.0-dot(_NormalTempXY,_NormalTempXY),0.0))));");
					break;
				}
			default:
				{
				}
			}

			//Recompile ALL Shaders
			CurrentSceneManager->GetRender()->RecompileShaders();

			//Create DS Shader Material
			jeStringc DsSource(DEFAULT_DS_SHADER);
			DsSource.ReplaceFirst("MAX_DS_BATCH_SIZE",jeStringc(MaxLightsPerPass).c_str());
			video::IShaderMaterial * DsSM = CurrentSceneManager->GetRender()->CreateShaderMaterial(DsSource);
			DeferredShadingMaterial->SetShaderMaterial(DsSM);
			DsSM->Drop();

			CameraNear = DeferredShadingMaterial->ShaderMaterial->GetVariable("Near");
			CameraFar = DeferredShadingMaterial->ShaderMaterial->GetVariable("Far");
			CameraFov = DeferredShadingMaterial->ShaderMaterial->GetVariable("Fov");
			CameraAspectRatio = DeferredShadingMaterial->ShaderMaterial->GetVariable("AspectRatio");
			LightsPosition = DeferredShadingMaterial->ShaderMaterial->GetVariable("PosLights");
			LightsCount = DeferredShadingMaterial->ShaderMaterial->GetVariable("LightsNum");

			//Get transparent global variables
			Transparent_LightPosition = CurrentSceneManager->GetRender()->GetGlobalShaderVariable("jeForward_TransparentLightPosition",true);
			Transparent_LightPosition->SetCachedSetup(false);
			Transparent_LightPosition->SetIgnoreVoidShaderVariables(true);
			Transparent_LightPosition->RemoveVoidShaderVariables();

			Transparent_LightPosition_Count = CurrentSceneManager->GetRender()->GetGlobalShaderVariable("jeLightsNumTransparent",true);
			Transparent_LightPosition_Count->SetCachedSetup(true);
			Transparent_LightPosition_Count->SetIgnoreVoidShaderVariables(true);
			Transparent_LightPosition_Count->RemoveVoidShaderVariables();

			//Create transparent setup callback
			JENEW(CurrentCallback,CTransparentLightPositionSetupCallback(Transparent_LightPosition,Transparent_LightPosition_Count))

			//Create Screen Quad Camera
			ScreenQuadCamera = CurrentSceneManager->CreateCamera();
			ScreenQuadCamera->SetCameraOptions(ICameraSceneNode::CameraOptions().Set2D());

		}

		//! Destructor
		CDeferredShadingCompositor::~CDeferredShadingCompositor()
		{
			CurrentCallback->Drop();
			Transparent_LightPosition_Count->Drop();
			Transparent_LightPosition->Drop();

			CurrentSceneManager->GetRoot()->RemoveChildren(ScreenQuadCamera);

			if(EnableAA)
			{
				FBO_AA->Drop();
				FBO_Downsample->Drop();
			}
			FBO_GBuffer->Drop();

			LightsCount->Drop();
			LightsPosition->Drop();
			CameraAspectRatio->Drop();
			CameraFov->Drop();
			CameraFar->Drop();
			CameraNear->Drop();

			CurrentSceneManager->GetRender()->UnSetShaderMaterialDefine("DS_GET_NORMAL_VALUE(__Normal,__TexCoord)");

			JEDELETE(DeferredShadingMaterial)
		}

		//! On Render
		void CDeferredShadingCompositor::OnRender()
		{
			//Recreate Textures if window size changes
			if(ScreenSize != CurrentSceneManager->GetRender()->GetContext()->GetWindow()->GetSize())
				CreateDSTextures(CurrentSceneManager->GetRender()->GetContext()->GetWindow()->GetSize());

			video::IRender * Render = CurrentSceneManager->GetRender();

			//first state, set gbufferfbo and create gbuffer
			jeMatrix4f CurrentViewMatrix;
			if(EnableAA)
			{
				Render->SetTarget(FBO_AA);
				Render->ClearScreen(video::jeColorARGB32f(0.0f,0.0f,0.0f,0.0f));
				CurrentSceneManager->GetCurrentCamera()->ApplyTransforms();
				CurrentViewMatrix = Render->GetMatrix(video::RMT_VIEW);
				CurrentSceneManager->RenderQueueGroups(false,RPG_SOLID);

				FBO_AA->CopyToBuffer(FBO_Downsample);
			}
			else
			{
				Render->SetTarget(FBO_GBuffer);
				Render->ClearDepth();
				CurrentSceneManager->GetCurrentCamera()->ApplyTransforms();
				CurrentViewMatrix = Render->GetMatrix(video::RMT_VIEW);
				CurrentSceneManager->RenderQueueGroups(false,RPG_SOLID);
			}

			//render to screen
			Render->SetTarget(NULL);
			Render->ClearScreen();

			ScreenQuadCamera->ApplyTransforms();

			jeMatrix4f ModelMatrix;
			ModelMatrix.Scale(jeVector3df((f32)ScreenSize.X,(f32)ScreenSize.Y,1));
			Render->SetMatrix(ModelMatrix,video::RMT_MODEL);

			//setup lights arrays for multipass rendering
			jeDynamicArray<CTransparentLightPositionSetupCallback::LightPositionsArrayType*> LightPositions;
			for(ISceneManager::LightListType::Iterator Iter = CurrentSceneManager->GetLightList()->begin();Iter!=CurrentSceneManager->GetLightList()->end();Iter++)
			{
				if((!LightPositions.GetSize()) || (LightPositions.GetLast()->GetSize() == MaxLightsPerPass))
				{
					CTransparentLightPositionSetupCallback::LightPositionsArrayType * Temp = NULL;
					JENEW(Temp,CTransparentLightPositionSetupCallback::LightPositionsArrayType(MaxLightsPerPass))
					LightPositions.InsertElement(Temp);
				}
				LightPositions.GetLast()->InsertElement(CurrentViewMatrix.MultiplyMatrixOnVector((*Iter)->GetPosition()));
			}

			//render solid objects
			Render->SetDepthMode();
			Render->SetMaterial(DeferredShadingMaterial);

			CameraNear->Setf32(CurrentSceneManager->GetCurrentCamera()->GetCameraOptions().CurrentNear);
			CameraFar->Setf32(CurrentSceneManager->GetCurrentCamera()->GetCameraOptions().CurrentFar);
			CameraFov->Setf32(CurrentSceneManager->GetCurrentCamera()->GetCameraOptions().CurrentFov*je::jeDegToRad/2.0f);
			CameraAspectRatio->Setf32((f32)ScreenSize.X / (f32)ScreenSize.Y);

			u1 NeedToRestoreOptions = false;

			for(u32 LightPassIndex=0;LightPassIndex<LightPositions.GetSize();LightPassIndex++)
			{
				LightsPosition->SetVector(LightPositions[LightPassIndex]->GetPtr(),LightPositions[LightPassIndex]->GetSize());
				LightsCount->Sets32(LightPositions[LightPassIndex]->GetSize());

				Render->RenderOperation(video::ROT_DRAW_QUAD);

				if(LightPassIndex < LightPositions.GetSize())
				{
					Render->SetDepthMode(video::DM_EQUAL);
					Render->SetGlobalMaterialType(true,video::MT_ADDITIVEBLEND);
					NeedToRestoreOptions = true;
				}
			}

			if(NeedToRestoreOptions)
			{
				Render->SetDepthMode();
				Render->SetGlobalMaterialType(false);
			}

			//render no light objects
			CurrentSceneManager->GetCurrentCamera()->ApplyTransforms();
			CurrentSceneManager->RenderQueueGroups(false,RPG_SOLID_WITHOUTLIGHTING);

			//render transparent objects
			Render->SetPreRenderOperationCallback(CurrentCallback);

			NeedToRestoreOptions = false;
			
			for(u32 i=0;i<CurrentSceneManager->GetCountOfObjectsInQueueGroup(scene::RPG_TRANSPARENT);i++)
			{
				if(NeedToRestoreOptions)
					Render->SetGlobalMaterialType(false);

				for(u32 LightPassIndex=0;LightPassIndex<LightPositions.GetSize();LightPassIndex++)
				{
					CurrentCallback->SetLightPositions(LightPositions[LightPassIndex]);
					CurrentSceneManager->RenderQueueObject(i,scene::RPG_TRANSPARENT);

					if(LightPassIndex < LightPositions.GetSize())
					{
						Render->SetGlobalMaterialType(true,video::MT_LIGHTBLEND);
						NeedToRestoreOptions = true;
					}
				}
			}

			Render->SetGlobalMaterialType(false);
			Render->SetPreRenderOperationCallback(NULL);

			//clear light arrays
			for(u32 i=0;i<LightPositions.GetSize();i++)
			{
				LightPositions[i]->Clear();
				JEDELETE(LightPositions[i])
			}
			LightPositions.Clear();

			ICompositor::OnRender();
		}

		//! Serialize
		void CDeferredShadingCompositor::Serialize(core::ISerializeStream * SerializeStream)
		{
			ICompositor::Serialize(SerializeStream);
		}

		//! Create Deferred Shading Textures
		s32 CDeferredShadingCompositor::CreateDSTextures(const jeVector2di & Size)
		{
			s32 Result = 0;

			ScreenSize = Size;

			video::ITexture * GBufferDiffuse = CurrentSceneManager->GetRender()->CreateTexture(ScreenSize,video::TF_RGBA8,video::TCF_DEFAULT);
			video::ITexture * GBufferNormal = CurrentSceneManager->GetRender()->CreateTexture(ScreenSize,video::TF_RG16F,video::TCF_DEFAULT);
			video::ITexture * GBufferDepth = CurrentSceneManager->GetRender()->CreateTexture(ScreenSize,video::TF_DEPTH32,video::TCF_DEFAULT);

			if(!FBO_GBuffer->AttachTexture(GBufferNormal,video::RTTL_COLOR0))
			{
				JELOG_E("GBuffer : Can't attach normal RG16F, change format to RGBA8")
				FBO_GBuffer->AttachTexture(NULL,video::RTTL_COLOR0);
				GBufferNormal->Drop();
				GBufferNormal = CurrentSceneManager->GetRender()->CreateTexture(ScreenSize,video::TF_RGBA8,video::TCF_DEFAULT);
				Result = 1;
				if(!FBO_GBuffer->AttachTexture(GBufferNormal,video::RTTL_COLOR0))
					JELOG_E("GBuffer : Can't attach normal RGBA8")
			}

			if(EnableAA)
			{
				s32 CurrentSamples = core::GetVariables()->GetVariablec8("SceneManager_DSCompositor_AA_Samples");
				if(CurrentSamples==-1)
					CurrentSamples = core::GetVariables()->GetVariablec8("Window_MSAA");

				video::ITexture * GBufferDiffuse_AA = CurrentSceneManager->GetRender()->CreateRenderBufferTexture(ScreenSize,video::TF_RGBA8,CurrentSamples);
				video::ITexture * GBufferNormal_AA = CurrentSceneManager->GetRender()->CreateRenderBufferTexture(ScreenSize,video::TF_RG16F,CurrentSamples);
				video::ITexture * GBufferDepth_AA = CurrentSceneManager->GetRender()->CreateRenderBufferTexture(ScreenSize,video::TF_DEPTH32,CurrentSamples);

				FBO_AA->AttachTexture(GBufferNormal_AA,video::RTTL_COLOR0);
				FBO_AA->AttachTexture(GBufferDiffuse_AA,video::RTTL_COLOR1);
				FBO_AA->AttachTexture(GBufferDepth_AA,video::RTTL_DEPTH);

				GBufferDepth_AA->Drop();
				GBufferNormal_AA->Drop();
				GBufferDiffuse_AA->Drop();

				FBO_Downsample->AttachTexture(GBufferNormal,video::RTTL_COLOR0);
				FBO_Downsample->AttachTexture(GBufferDiffuse,video::RTTL_COLOR1);
				FBO_Downsample->AttachTexture(GBufferDepth,video::RTTL_DEPTH);
			}
			else
			{
				if(!FBO_GBuffer->AttachTexture(GBufferDiffuse,video::RTTL_COLOR1))
					JELOG_E("GBuffer : Can't attach diffuse RGBA8")
			}

			if(!FBO_GBuffer->AttachTexture(GBufferDepth,video::RTTL_DEPTH))
				JELOG_E("GBuffer : Can't attach DEPTH32F")

			DeferredShadingMaterial->SetTexture(GBufferDiffuse,0);
			DeferredShadingMaterial->SetTexture(GBufferNormal,1);
			DeferredShadingMaterial->SetTexture(GBufferDepth,2);

			GBufferDiffuse->Drop();
			GBufferNormal->Drop();
			GBufferDepth->Drop();

			return Result;
		}
	}
}
